
/// <reference types="@webgpu/types"/>

import { GUI } from 'dat.gui'

import { mat4 } from 'wgpu-matrix'

import kleeImg from '../assets/klee.jpg'
import shaderSource from '../assets/shader.wgsl?raw'
import shaderSourceOL from '../assets/shader_outline.wgsl?raw'
import { Tessellator } from './tessellator'


export function canRunWebGPU() {
    return Boolean(navigator.gpu)
}
const degToRad = (d: number) => d * Math.PI / 180;

const settings = {
    rotation: [0, 20, degToRad(0)],
    translation: [0, 0, 3],
    fov: 30
};
const gui = new GUI();
gui.addFolder('rotation controls')
gui.add(settings.rotation, 0, -180, 180, 1).name('rotation.x');
gui.add(settings.rotation, 1, -180, 180, 1).name('rotation.y');
gui.add(settings.rotation, 2, -180, 180, 1).name('rotation.z');

gui.addFolder('translation controls')
gui.add(settings.translation, 0, -5, 5, 0.1).name('translate.x');
gui.add(settings.translation, 1, -5, 5, 0.1).name('translate.y');
gui.add(settings.translation, 2, -5, 5, 0.1).name('translate.z');

gui.addFolder('projection controls')
gui.add(settings, 'fov', 5, 170, 1).name('projection.fov');

const tessellator = new Tessellator()

tessellator.addFace([-.5, -.5, .5], [1, 0, 0], [0, 1, 0])
tessellator.addFace([.5, -.5, -.5], [-1, 0, 0], [0, 1, 0])
tessellator.addFace([-.5, -.5, -.5], [0, 0, 1], [0, 1, 0])
tessellator.addFace([.5, -.5, .5], [0, 0, -1], [0, 1, 0])
tessellator.addFace([-.5, .5, .5], [1, 0, 0], [0, 0, -1])
tessellator.addFace([.5, -.5, -.5], [0, 0, 1], [-1, 0, 0])

const vertices = new Float32Array(tessellator.vectices)
const uvs = new Float32Array(tessellator.uvs)
const normals = new Float32Array(tessellator.normals)

async function loadImageBitmap() {
    const blob = await (await fetch(kleeImg)).blob()

    return await createImageBitmap(blob, { colorSpaceConversion: 'none' })
}

export async function main(canvas: HTMLCanvasElement) {
    const context = <GPUCanvasContext>canvas.getContext('webgpu')

    const adaptor = await navigator.gpu.requestAdapter()
    const device = await adaptor?.requestDevice()

    if (device === undefined)
        return

    const canvasFormat = navigator.gpu.getPreferredCanvasFormat()

    context.configure({
        device: device,
        format: canvasFormat
    })

    const shaderModule = device.createShaderModule({
        label: 'my shader',
        code: shaderSource
    })

    const shaderModuleOL = device.createShaderModule({
        label: 'my shader-outline',
        code: shaderSourceOL
    })

    let depthTexture = device.createTexture({
        label: 'depth',
        format: 'depth32float',
        usage: GPUTextureUsage.RENDER_ATTACHMENT,
        size: [canvas.width, canvas.height]
    })

    const descriptor: GPURenderPassDescriptor = {
        colorAttachments: [
            {
                clearValue: [0.3, 0.3, 0.3, 1],
                loadOp: 'clear',
                storeOp: 'store',
                view: context.getCurrentTexture().createView()
            }
        ],

        depthStencilAttachment: {
            view: depthTexture.createView(),

            depthClearValue: 1.0,
            depthLoadOp: 'clear',
            depthStoreOp: 'store'
        },
    }

    const vertexBuffer = device.createBuffer({
        label: 'vertices buffer',
        size: vertices.byteLength,
        usage: GPUBufferUsage.VERTEX | GPUBufferUsage.COPY_DST
    })
    device.queue.writeBuffer(vertexBuffer, 0, vertices)

    const uvBuffer = device.createBuffer({
        label: 'uv buffer',
        size: uvs.byteLength,
        usage: GPUBufferUsage.VERTEX | GPUBufferUsage.COPY_DST
    })
    device.queue.writeBuffer(uvBuffer, 0, uvs)

    const normalBuffer = device.createBuffer({
        label: 'normal buffer',
        size: normals.byteLength,
        usage: GPUBufferUsage.VERTEX | GPUBufferUsage.COPY_DST
    })
    device.queue.writeBuffer(normalBuffer, 0, normals)

    // ========= Buffer Layout Configuration ========= {
    const vertexBufferLayuout: GPUVertexBufferLayout = {
        arrayStride: 4 * 4,
        attributes: [{
            format: 'float32x4',
            offset: 0,
            shaderLocation: 0
        }]
    }

    const uvBufferLayuout: GPUVertexBufferLayout = {
        arrayStride: 2 * 4,
        attributes: [{
            format: 'float32x2',
            offset: 0,
            shaderLocation: 1
        }]
    }
    
    const normalBufferLayuout: GPUVertexBufferLayout = {
        arrayStride: 4 * 4,
        attributes: [{
            format: 'float32x4',
            offset: 0,
            shaderLocation: 2
        }]
    }

    // ========== Pipeline Creating ========== {

    const pipeline = device.createRenderPipeline({
        label: 'my render pipeline',
        layout: 'auto',
        vertex: {
            module: shaderModule,
            entryPoint: 'vert',
            buffers: [ 
                vertexBufferLayuout,
                uvBufferLayuout,
                normalBufferLayuout,
            ]
        },
        fragment: {
            module: shaderModule,
            entryPoint: 'frag',
            targets: [{ format: canvasFormat }]
        },
        depthStencil: {
            depthWriteEnabled: true,
            depthCompare: 'less',
            format: 'depth32float'
        },
        primitive: {
            cullMode: 'front'
        }
    })

    const pipelineOutline = device.createRenderPipeline({
        label: 'outline render pipeline',
        layout: 'auto',
        vertex: {
            module: shaderModuleOL,
            entryPoint: 'vert',
            buffers: [
                vertexBufferLayuout,
            ]
        },
        fragment: {
            module: shaderModuleOL,
            entryPoint: 'frag',
            targets: [{ format: canvasFormat }]
        },
        depthStencil: {
            depthWriteEnabled: false,
            depthCompare: 'less',
            format: 'depth32float'
        },
        primitive: {
            cullMode: 'front'
        }
    })

    // } ========== Pipeline Creating ==========
    
    // } ========= Buffer Layout Configuration =========

    const ofsUniformArray = new Float32Array([0, 0, 0])

    const ofsUniformBuffer = device.createBuffer({
        label: 'myUniformBuffer',
        size: ofsUniformArray.byteLength,
        usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST
    })

    const matUniformBuffer = device.createBuffer({
        label: 'mat uniform',
        size: 16 * 4 * 2,
        usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST
    })

    // ========= Texture Loading ========== {

    const texBitmap = await loadImageBitmap()
    const textureBuffer = device.createTexture({
        label: 'image texture',
        format: 'rgba8unorm',
        size: [texBitmap.width, texBitmap.height],  // [width, height]
        usage: GPUTextureUsage.TEXTURE_BINDING |
               GPUTextureUsage.COPY_DST |  // required
               GPUTextureUsage.RENDER_ATTACHMENT  // required, for: copyExternalImageToTexture
    })

    device.queue.copyExternalImageToTexture(
        { source: texBitmap, flipY: false },
        { texture: textureBuffer },
        { width: texBitmap.width, height: texBitmap.height }
    )

    const texSampler = device.createSampler({})
    
    // } ========= Texture Loading ==========

    const bindGroup = device.createBindGroup({
        label: 'group',
        layout: pipeline.getBindGroupLayout(0),
        entries: [
            { 
                binding: 1, 
                resource: textureBuffer.createView()
            },

            {
                binding: 2,
                resource: texSampler
            },
            
            {
                binding: 3,
                resource: { buffer: matUniformBuffer }
            }
        ]
    })

    const bindGroupOL = device.createBindGroup({
        label: 'group-outline',
        layout: pipelineOutline.getBindGroupLayout(0),
        entries: [
            {
                binding: 0,
                resource: { buffer: matUniformBuffer }
            },
            {
                binding: 1,
                resource: { buffer: ofsUniformBuffer }
            }
        ]
    })

    function render(dt: number) {
        if (device == null)
            return

        const [ctxW, ctxH] = [context.canvas.width, context.canvas.height]

        // 上一个 Canvas 的 texture 生命周期在 request animation frame 时候就结束了。
        descriptor.colorAttachments[0].view = context.getCurrentTexture().createView()

        if (depthTexture.width != ctxW || depthTexture.height != ctxH) {
            depthTexture.destroy()

            depthTexture = device.createTexture({
                size: [ctxW, ctxH],
                usage: GPUTextureUsage.RENDER_ATTACHMENT,
                format: 'depth32float'
            })
            descriptor.depthStencilAttachment.view = depthTexture.createView()
        }

        ofsUniformArray[0] = Math.sin(degToRad(dt + 70)) * 0.5 + 0.5
        ofsUniformArray[1] = Math.sin(degToRad(dt) * 0.5) * 0.5 + 0.5
        ofsUniformArray[2] = Math.sin(degToRad(dt) * 0.8) * 0.5 + 0.5
        device.queue.writeBuffer(ofsUniformBuffer, 0, ofsUniformArray)

        const matRot = mat4.identity()
        mat4.rotateZ(matRot, degToRad(settings.rotation[2]), matRot)
        mat4.rotateY(matRot, degToRad(settings.rotation[1]), matRot)
        mat4.rotateX(matRot, degToRad(settings.rotation[0]), matRot)
        const matTsl = mat4.translation(settings.translation)

        const matM = mat4.multiply(matTsl, matRot)
        const matP = mat4.perspective(degToRad(settings.fov), context.canvas.width / context.canvas.height, 0.1, 100)
        const matV = mat4.inverse(mat4.lookAt([0, 0, -1], [0, 0, 0], [0, 1, 0]))

        const matMVP = mat4.create()
        const matMV = mat4.create()
        mat4.multiply(matV, matM, matMV)
        mat4.multiply(matP, matMV, matMVP)

        device.queue.writeBuffer(matUniformBuffer, 0, matMVP)
        device.queue.writeBuffer(matUniformBuffer, 16 * 4, matM)

        const encoder = device.createCommandEncoder({label: 'my encoder'})
        const pass = encoder.beginRenderPass(descriptor)

        pass.setVertexBuffer(0, vertexBuffer)
        pass.setVertexBuffer(1, uvBuffer)
        pass.setVertexBuffer(2, normalBuffer)

        pass.setPipeline(pipelineOutline)
        pass.setBindGroup(0, bindGroupOL)
        pass.draw(vertices.length / 4)

        pass.setPipeline(pipeline)
        pass.setBindGroup(0, bindGroup)
        pass.draw(vertices.length / 4)

        pass.end()

        device.queue.submit([ encoder.finish() ])

        requestAnimationFrame(render)
    }

    render()
}
